import radar_read as radar
import water_read as precipitation
import numpy as np
import os
import config
import numba
# from skimage.transform import resize
import time


def print_matrix(data):
    for i in range(data.shape[0]):
        for k in range(data.shape[1]):
            print(data[i][k], end=' ')
        print("")


# 缩放矩阵,有效率问题
@numba.jit
def resize(src, dstsize):  # 输入src 和size
    # if src.ndim == 3:
    #     dstsize.append(3)
    dst = np.array(np.zeros(dstsize), src.dtype)
    factory = float(np.size(src, 0)) / dstsize[0]
    factorx = float(np.size(src, 1)) / dstsize[1]
    # print('factory', factory, 'factorx', factorx)
    srcheight = np.size(src, 0)
    srcwidth = np.size(src, 1)
    # print('srcwidth', srcwidth, 'srcheight', srcheight)
    for i in range(dstsize[0]):
        for j in range(dstsize[1]):
            y = float(i) * factory
            x = float(j) * factorx
            if y + 1 > srcheight:  # 越界判断
                y -= 1
            if x + 1 > srcwidth:
                x -= 1

            dst[i, j] = src[int(y), int(x)]
    return dst


def grib_pre_process(rdata, grb):
    radar_lon_lat = radar.lon_lat_bounds(rdata)
    grb_lon_lat = precipitation.lon_lat_bounds(grb)

    # 计算交集经纬度
    lon_min = max(radar_lon_lat[0][0], grb_lon_lat[0][0])
    lon_max = min(radar_lon_lat[0][1], grb_lon_lat[0][1])
    # print("lon_min:", lon_min, "lon_max", lon_max)
    lat_min = max(radar_lon_lat[1][0], grb_lon_lat[1][0])
    lat_max = min(radar_lon_lat[1][1], grb_lon_lat[1][1])
    # print("lat_min:", lat_min, "lat_max", lat_max)

    # 计算降水矩阵边界
    lat, lon = grb.latlons()

    # 经度区域索引
    lon_range = lon[0]
    # for x in lon_range:
    #     print(x, end=' ')
    # print("")
    lon_start = np.where(lon_range >= lon_min)[0][0] - 1
    # print("lon_start,index:", lon_start, "value:", lon_range[lon_start])

    lon_end = np.where(lon_range <= lon_max)
    lon_end = lon_end[0][lon_end[0].max()]
    # print("lon_end,index:", lon_end, "value:", lon_range[lon_end])

    # 纬度区域索引
    lat_range = lat[:, :1].reshape(lat.shape[0])
    # for x in lat_range:
    #     print(x, end=' ')
    # print("")

    lat_start = np.where(lat_range >= lat_min)[0][0]
    # print("lat_start,index:", lat_start, "value:", lat_range[lat_start])

    lat_end = np.where(lat_range <= lat_max)
    lat_end = lat_end[0][lat_end[0].max()]
    # print("lat_end,index:", lat_end, "value:", lat_range[lat_end])

    # 裁切交集区域数据
    lon = lon[lat_start:lat_end + 1, lon_start:lon_end + 1]
    lat = lat[lat_start:lat_end + 1, lon_start:lon_end + 1]
    data = grb.values
    data = data[lat_start:lat_end + 1, lon_start:lon_end + 1]

    lat = lat[::-1]
    # print_matrix(lat)

    data = data[::-1]
    data = resize(data, rdata['data'].shape)

    # precipitation.show_map(lon, lat)
    # precipitation.show_hot(sdata)

    return lat, lon, data


def save_grib_png():
    radar_file = 'CREF/ACHN_CREF_20171113_035000.bin'
    radarInfo = radar.readFile(radar_file)  # 读取雷达文件

    files = os.listdir("SURF_CMPA_FRT_5KM_10MIN")
    for f in files:
        precipitation_file = "SURF_CMPA_FRT_5KM_10MIN/" + f
        grb = precipitation.read_file(precipitation_file)  # 读取降水文件
        lat, lon, data = grib_pre_process(radarInfo, grb)
        precipitation.save_fig(data, "out/surf/" + f + ".png")


# 降水数据预处理为(72,4200,6200)的数组
def save_grib_to_np():
    radar_file = config.CREF_FILES + 'ACHN_CREF_20171113_035000.bin'
    radarInfo = radar.readFile(radar_file)  # 读取雷达文件
    files = os.listdir(config.SURF_CMPA_FRT_5KM_10MIN)
    files.sort(key=lambda x: int(x[-9:-5]))

    size = len(files)
    grib_array = np.zeros((size, 4200, 6200), dtype=np.float64)
    print("create array,shape:", grib_array.shape)

    start = time.time()
    i = 0
    for f in files:
        precipitation_file = config.SURF_CMPA_FRT_5KM_10MIN + f
        print("process file:", precipitation_file)
        grb = precipitation.read_file(precipitation_file)  # 读取降水文件

        print("datetime:", grb['year'], grb['month'], grb['day'], grb['hour'], grb['minute'], grb['second'])

        lat, lon, data = grib_pre_process(radarInfo, grb)
        # precipitation.save_fig(data, "out/surf/" + f + ".png")
        grib_array[i] = data
        i += 1
    grib_array[grib_array==9999] = 0
    end = time.time()
    processTime = end - start
    print("process time:", processTime)
    np.save(config.WATER_FILE, grib_array)


def save_radar_to_np():
    files = os.listdir(config.CREF_FILES)
    files.sort(key=lambda x: int(x[19:-4]))

    size = len(files)
    radar_array = np.zeros((size, 4200, 6200), dtype=np.float64)
    print("create array,shape:", radar_array.shape)

    start = time.time()
    i = 0
    for f in files:
        radar_file = config.CREF_FILES + f
        print("process file:", radar_file)
        radarInfo = radar.readFile(radar_file)
        headInfo = radarInfo['headInfo']
        radar_array[i] = radarInfo['data']
        print("datetime:", headInfo['year'], headInfo['month'], headInfo['day'], headInfo['hour'], headInfo['minite'],
              headInfo['second'], )
        i += 1
    radar_array = np.clip(radar_array, 0, None)
    radar_array /= 1000     # fanwei [0, 1]
    end = time.time()
    processTime = end - start
    print("process time:", processTime)

    np.save(config.RADAR_FILE, radar_array)


if __name__ == '__main__':
    # radar_file = 'CREF/ACHN_CREF_20171113_035000.bin'
    # precipitation_file = 'SURF_CMPA_FRT_5KM_10MIN/Z_SURF_C_BABJ_20171113001802_P_CMPA_FAST_CHN_0P05_10MIN-PRE-201711130010.GRB2'
    #
    # rdata = radar.readFile(radar_file)  # 读取雷达文件
    # grb = precipitation.read_file(precipitation_file)  # 读取降水文件
    #
    # lat, lon, data = grib_pre_process(rdata, grb)
    # precipitation.show_hot(data)

    # save_grib_to_np()
    save_radar_to_np()
